#ifndef GEOLOGICALHAZARD_H
#define GEOLOGICALHAZARD_H
#pragma once

#ifdef _WIN32
#include <Windows.h>
#include <direct.h>
#include <io.h>
#endif

#include <string>
#include <vector>
#include <cstdio>
#include <opencv2/opencv.hpp>
#include "onnxruntime_cxx_api.h"

#ifdef USE_CUDA
#include <cuda_fp16.h>
#endif

typedef struct  SegmentContentType
{
    cv::Mat class_mat; // 21|? classes 1 channel
    cv::Mat color_mat; // 21 colors different classes, 3 channels.
    std::unordered_map<int, std::string> names_map;
    bool flag;

    SegmentContentType() : flag(false)
    {};
} SegmentContent;
enum
{
    CHW = 0, HWC = 1
};
class GeologicalHazardSeg
{
    private:
        Ort::Env ort_env;
        Ort::Session *ort_session = nullptr;
        std::vector<const char *> input_node_names;
        std::vector<std::string> input_node_names_;
        std::vector<std::vector<int64_t>> dynamic_input_node_dims; // >=1 inputs.
        unsigned int dynamic_input_height; // init only, will change according to input mat.
        unsigned int dynamic_input_width; // init only, will change according to input mat.
        unsigned int dynamic_input_tensor_size = 1; // init only, will change according to input mat.
        Ort::MemoryInfo memory_info_handler = Ort::MemoryInfo::CreateCpu(
            OrtArenaAllocator, OrtMemTypeDefault);
        std::vector<const char *> output_node_names;
        std::vector<std::string> output_node_names_;
        const char *log_id = nullptr;
        unsigned int num_outputs = 1;
        unsigned int num_inputs = 1;
        std::vector<float> dynamic_input_values_handler;

   protected:
        const unsigned int num_threads; // initialize at runtime.

    public:
        // single input with dynamic height and width.
        explicit GeologicalHazardSeg(const std::string &_onnx_path, unsigned int _num_threads, int height, int width);

        ~GeologicalHazardSeg();

    protected:
        GeologicalHazardSeg(const GeologicalHazardSeg&) = delete;

        GeologicalHazardSeg(GeologicalHazardSeg&&) = delete;

        GeologicalHazardSeg&operator=(const GeologicalHazardSeg&) = delete;

        GeologicalHazardSeg&operator=(GeologicalHazardSeg&&) = delete;

    private:
        const float mean_vals[3] = {0.485f, 0.456f, 0.406f};
        const float scale_vals[3] = {1.f / 0.229f, 1.f / 0.224f, 1.f / 0.225f};
        const char *class_names[1] = {
            "aeroplane"}; // 1 classes

    private:
        Ort::Value transform(const cv::Mat &mat);

        void print_debug_string();

        std::string OrtCompatiableGetInputName(size_t index, OrtAllocator* allocator,
            Ort::Session* ort_session);
        std::string OrtCompatiableGetOutputName(size_t index, OrtAllocator* allocator,
            Ort::Session* ort_session);
        Ort::Value create_tensor(const cv::Mat& mat, const std::vector<int64_t>& tensor_dims,
            const Ort::MemoryInfo& memory_info_handler,
            std::vector<float>& tensor_value_handler,
            unsigned int data_format);
        cv::Mat normalize(const cv::Mat& mat, float mean, float scale);
        cv::Mat normalize(const cv::Mat& mat, const float* mean, const float* scale);
        void normalize(const cv::Mat& inmat, cv::Mat& outmat,
            float mean, float scale);
        void normalize_inplace(cv::Mat& mat_inplace, float mean, float scale);
        void normalize_inplace(cv::Mat& mat_inplace, const float* mean, const float* scale);

    public:
        void detect(const cv::Mat &mat, SegmentContent &content);

};

#endif 
